package cn.iocoder.boot.springai001.service;

import cn.iocoder.boot.springai001.utils.MediaUtils;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;

/**
 * 流式工具调用服务
 * <p>
 * 基于Spring AI官方推荐的手动工具调用模式实现流式输出
 */
@Slf4j
@Service
public class StreamingToolCallingService {

    @Resource
    private ChatClient chatClient;
    @Resource
    private ChatMemory chatMemory;
    @Resource
    private ToolCallingManager toolCallingManager;

    public Flux<ChatResponse> streamingToolChat(String systemPrompt, String query, List<String> imageUrls,
                                                String model, Double temperature, String conversationId) {

        ChatOptions chatOptions = ToolCallingChatOptions.builder()
                .model(model)
                .internalToolExecutionEnabled(false)
                .temperature(temperature)
                .build();

        Prompt prompt = createPrompt(systemPrompt, query, imageUrls, chatOptions, conversationId);

        return streamingToolChatRecursive(prompt, chatOptions, conversationId);
    }

    private Prompt createPrompt(String systemPrompt, String query, List<String> imageUrls,
                                ChatOptions chatOptions, String conversationId) {
        List<Message> messagesAll = chatMemory.get(conversationId);
        messagesAll.add(new SystemMessage(systemPrompt));
        messagesAll.add(UserMessage.builder().text(query).media(MediaUtils.mediaArray(imageUrls)).build());
        return Prompt.builder().messages(messagesAll).chatOptions(chatOptions).build();
    }

    private Flux<ChatResponse> streamingToolChatRecursive(Prompt prompt, ChatOptions chatOptions, String conversationId) {
        return chatClient.prompt(prompt)
                .options(chatOptions)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
                .stream()
                .chatResponse()
                .concatMap(chatResponse -> {
                    if (chatResponse.hasToolCalls()) {
                        // 1. 执行工具调用
                        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, chatResponse);

                        // 2. 返回工具执行结果的 ChatResponse
                        ChatResponse toolResultResponse = createToolResultChatResponse(toolExecutionResult);

                        // 3. 获取工具执行结果消息
                        List<Message> messages = toolExecutionResult.conversationHistory();

                        // 4. 创建新的提示继续AI对话
                        Prompt nextPrompt = new Prompt(messages, chatOptions);

                        // 5. 清空会话内存
                        chatMemory.clear(conversationId);

                        // 6. 返回：工具调用信息 + 工具执行结果 + 递归的AI响应
                        return Flux.just(chatResponse, toolResultResponse)
                                .concatWith(streamingToolChatRecursive(nextPrompt, chatOptions, conversationId));
                    } else {
                        // 直接返回AI响应
                        return Flux.just(chatResponse);
                    }
                }).doOnError(error -> log.error("对话处理出错，会话ID: {}, 错误: {}",
                        conversationId, error.getMessage(), error))
                .doOnComplete(() ->
                        log.info("会话结束，会话ID: {}", conversationId)
                );
    }

    private ChatResponse createToolResultChatResponse(ToolExecutionResult toolExecutionResult) {
        // 倒数第一个是工具执行结果
        Message message = extractToolResultMessage(toolExecutionResult);
        ToolResponseMessage.ToolResponse toolResponse = ((ToolResponseMessage) message).getResponses().get(0);
        Generation generation = new Generation(new AssistantMessage(toolResponse.responseData(),
                Map.of("id", toolResponse.id(), "name", toolResponse.name(), "type", "function")),
                ChatGenerationMetadata.builder()
                        .finishReason("TOOL")
                        .build());
        return ChatResponse.builder()
                .generations(List.of(generation))
                .build();
    }

    /**
     * 从ToolExecutionResult中提取工具执行结果消息
     */
    private Message extractToolResultMessage(ToolExecutionResult toolExecutionResult) {
        // 获取最后一个消息，应该是工具执行结果
        List<Message> conversationHistory = toolExecutionResult.conversationHistory();
        return conversationHistory.get(conversationHistory.size() - 1);
    }

}
